import os
import cv2
import numpy as np
import matplotlib.pyplot as plt

import matplotlib

matplotlib.use("agg")

current_dir = os.path.dirname(os.path.abspath(__file__))


# -----------------histogram-------------------
def calculate_pixel_frequency(image):
    """
    计算图像的像素频率（直方图）
    参数：
        image: 输入的图像（numpy array）
    返回：
        包含每个灰度级像素数量的数组
    """
    pixel_frequency = np.zeros(256, dtype=int)
    for pixel_value in image.flatten():
        pixel_frequency[pixel_value] += 1
    return pixel_frequency


def perform_histogram_equalization(image):
    """
    对图像进行直方图均衡化
    参数：
        image: 输入的图像（numpy array）
    返回：
        均衡化后的图像
    """
    histogram = calculate_pixel_frequency(image)
    cdf = histogram.cumsum()
    cdf_normalized = ((cdf - cdf.min()) * 255 / (cdf.max() - cdf.min())).astype("uint8")
    equalized_image = cdf_normalized[image]
    return equalized_image


def save_histogram_plot(histogram, title, output_filename):
    """
    保存直方图图像
    参数：
        histogram: 像素频率数组
        title: 图像标题
        output_filename: 输出文件名
    """
    plt.figure()
    plt.title(title)
    plt.xlabel("Pixel Value")
    plt.ylabel("Frequency")
    plt.bar(range(256), histogram, width=1, color="black")
    plt.xlim([0, 255])

    plt.savefig(output_filename)
    plt.close()


def enhance_input_image(image_path):
    """
    增强输入图像
    参数：
        image_path: 输入图像路径
    返回：
        增强后的图像
    """
    # 读取图像
    input_image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

    # 计算并保存原始图像的直方图
    original_histogram = calculate_pixel_frequency(input_image)
    base_name = os.path.basename(image_path)
    image_name, extension = os.path.splitext(base_name)
    original_histogram_filename = os.path.join(
        current_dir + "/test1/output", f"{image_name}_original_hist.jpg"
    )
    save_histogram_plot(
        original_histogram, "Original Image Histogram", original_histogram_filename
    )

    # 直方图均衡化
    enhanced_image = perform_histogram_equalization(input_image)

    # 计算并保存增强后图像的直方图
    enhanced_histogram = calculate_pixel_frequency(enhanced_image)
    enhanced_histogram_filename = os.path.join(
        current_dir + "/test1/output", f"{image_name}_enhanced_hist.jpg"
    )
    save_histogram_plot(
        enhanced_histogram, "Enhanced Image Histogram", enhanced_histogram_filename
    )

    # 获取增强后文件名并保存增强后的图像
    enhanced_file_name = f"{image_name}_enhanced.jpg"
    enhanced_file_path = os.path.join(current_dir + "/test1/output", enhanced_file_name)
    cv2.imwrite(enhanced_file_path, enhanced_image)
    print(f"Enhanced image saved to {enhanced_file_path}")

    # 创建四合一图像
    fig, axs = plt.subplots(2, 2, figsize=(10, 8))
    axs[0, 0].imshow(cv2.imread(original_histogram_filename), cmap="gray")
    axs[0, 0].axis("off")

    axs[0, 1].imshow(cv2.imread(enhanced_histogram_filename), cmap="gray")
    axs[0, 1].axis("off")

    axs[1, 0].imshow(cv2.imread(image_path))
    axs[1, 0].set_title("Original Image")
    axs[1, 0].axis("off")

    axs[1, 1].imshow(cv2.imread(enhanced_file_path), cmap="gray")
    axs[1, 1].set_title("Enhanced Image")
    axs[1, 1].axis("off")
    plt.tight_layout()
    combined_filename = os.path.join(
        current_dir + "/test1/output", "combined_result.jpg"
    )
    plt.savefig(combined_filename)
    _, img_encoded = cv2.imencode(".png", cv2.imread(combined_filename))
    return img_encoded.tobytes()


def brightness_contrast(image, alpha=1.2, beta=30):  # 调整亮度和对比度
    new_image = np.clip(image * alpha + beta, 0, 255).astype(
        np.uint8
    )  # alpha控制对比度，beta控制亮度；使用np.clip将调整后的值限制在0到255之间，确保像素值有效
    return new_image


def sharpen(image):  # 自定义锐化滤波器
    height, width, channels = image.shape
    kernel = np.array(
        [
            [0, -0.1, 0],  # 定义锐化核：使用一个3x3卷积核来增强图像细节
            [-0.1, 2, -0.1],
            [0, -0.1, 0],
        ]
    )

    # 创建新图像
    sharpened_image = np.zeros_like(
        image
    )  # 初始化一个与原图像相同大小的空数组用于存储锐化后的图像。

    # 遍历每个像素，忽略边缘，避免越界
    for y in range(1, height - 1):
        for x in range(1, width - 1):
            for c in range(channels):
                # 计算卷积
                value = (  # 对邻域像素和卷积核元素相乘并求和
                    kernel[0, 0] * image[y - 1, x - 1, c]
                    + kernel[0, 1] * image[y - 1, x, c]
                    + kernel[0, 2] * image[y - 1, x + 1, c]
                    + kernel[1, 0] * image[y, x - 1, c]
                    + kernel[1, 1] * image[y, x, c]
                    + kernel[1, 2] * image[y, x + 1, c]
                    + kernel[2, 0] * image[y + 1, x - 1, c]
                    + kernel[2, 1] * image[y + 1, x, c]
                    + kernel[2, 2] * image[y + 1, x + 1, c]
                )
                # 限制值范围
                sharpened_image[y, x, c] = np.clip(value, 0, 255)

    return sharpened_image


def enhance_brightness_and_sharpen(image_path, output_dir, alpha=1.2, beta=30):
    # 读取彩色图像
    img = cv2.imread(image_path)

    # 增强亮度
    bright_img = brightness_contrast(img, alpha=alpha, beta=beta)

    # 应用锐化滤波器
    sharpened_img = sharpen(bright_img)

    # 获取原始文件名并创建新的文件名
    base_name = os.path.basename(image_path)
    name = os.path.splitext(base_name)
    enhanced_file_name = f"{name[0]}_enhanced.jpg"
    enhanced_file_path = os.path.join(output_dir, enhanced_file_name)

    # 保存增强后的图像
    cv2.imwrite(enhanced_file_path, sharpened_img)
    print(f"Enhanced image saved to {enhanced_file_path}")

    # 显示原始图像和增强后的图像
    plt.figure(figsize=(10, 4))

    plt.subplot(1, 2, 1)
    plt.title("Original Image")
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.axis("off")

    plt.subplot(1, 2, 2)
    plt.title("Enhanced Image")
    plt.imshow(cv2.cvtColor(sharpened_img, cv2.COLOR_BGR2RGB))
    plt.axis("off")
    plt.savefig("test1/output/result.jpg")
    plt.show()
    return sharpened_img


# -----------------linear transformation-------------------
def brightness_contrast(image, alpha=1.2, beta=30):  # 调整亮度和对比度
    new_image = np.clip(image * alpha + beta, 0, 255).astype(
        np.uint8
    )  # alpha控制对比度，beta控制亮度；使用np.clip将调整后的值限制在0到255之间，确保像素值有效
    return new_image


def sharpen(image):  # 自定义锐化滤波器
    height, width, channels = image.shape
    kernel = np.array(
        [
            [0, -0.1, 0],  # 定义锐化核：使用一个3x3卷积核来增强图像细节
            [-0.1, 2, -0.1],
            [0, -0.1, 0],
        ]
    )

    # 创建新图像
    sharpened_image = np.zeros_like(
        image
    )  # 初始化一个与原图像相同大小的空数组用于存储锐化后的图像。

    # 遍历每个像素，忽略边缘，避免越界
    for y in range(1, height - 1):
        for x in range(1, width - 1):
            for c in range(channels):
                # 计算卷积
                value = (  # 对邻域像素和卷积核元素相乘并求和
                    kernel[0, 0] * image[y - 1, x - 1, c]
                    + kernel[0, 1] * image[y - 1, x, c]
                    + kernel[0, 2] * image[y - 1, x + 1, c]
                    + kernel[1, 0] * image[y, x - 1, c]
                    + kernel[1, 1] * image[y, x, c]
                    + kernel[1, 2] * image[y, x + 1, c]
                    + kernel[2, 0] * image[y + 1, x - 1, c]
                    + kernel[2, 1] * image[y + 1, x, c]
                    + kernel[2, 2] * image[y + 1, x + 1, c]
                )
                # 限制值范围
                sharpened_image[y, x, c] = np.clip(value, 0, 255)

    return sharpened_image


def enhance_brightness_and_sharpen(image_path, output_dir, alpha=1.2, beta=30):
    # 读取彩色图像
    img = cv2.imread(image_path)

    # 增强亮度
    bright_img = brightness_contrast(img, alpha=alpha, beta=beta)

    # 应用锐化滤波器
    sharpened_img = sharpen(bright_img)

    # 获取原始文件名并创建新的文件名
    base_name = os.path.basename(image_path)
    name = os.path.splitext(base_name)
    enhanced_file_name = f"{name[0]}_enhanced.jpg"
    enhanced_file_path = os.path.join(output_dir, enhanced_file_name)

    # 保存增强后的图像
    cv2.imwrite(enhanced_file_path, sharpened_img)
    print(f"Enhanced image saved to {enhanced_file_path}")

    # 显示原始图像和增强后的图像
    plt.figure(figsize=(10, 4))

    plt.subplot(1, 2, 1)
    plt.title("Original Image")
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.axis("off")

    plt.subplot(1, 2, 2)
    plt.title("Enhanced Image")
    plt.imshow(cv2.cvtColor(sharpened_img, cv2.COLOR_BGR2RGB))
    plt.axis("off")
    _, img_encoded = cv2.imencode(".png", cv2.imread(enhanced_file_path))
    return img_encoded.tobytes()
